import torch
from PIL import Image
from torchvision import transforms as T

""" 对比两个API的卷积 """
img = Image.open("../../images/lena.png")
transformer = T.Compose([
    T.Resize((100, 100), antialias=True),
    T.ToTensor()
])
img = transformer(img)
print(img.shape)    # torch.Size([3, 100, 100]) CHW

# ———————调用函数———————————————————————————————————————————————————————————
# 输入3通道 输出 1通道
'''
通过上个文件的分析，我们知道
卷积核的维度 = (输出通道数, 输入通道数, 高度, 宽度)
kernel = (1, 3, 3, 3)
'''
# kernel = torch.tensor([
#     [
#         [
#             [-1, 0, 1],
#             [-2, 0, 2],
#             [-1, 0, 1]
#         ],
#         [
#             [-1, 0, 1],
#             [-1, 0, 1],
#             [-2, 0, 2]
#         ],
#         [
#             [1, 0, -1],
#             [2, -2, 0],
#             [1, -1, 0]
#         ]
#     ]
# ]).float()
# kernel = torch.rand((1, 3, 3, 3))
kernel = torch.rand((3, 3, 5, 5))
print(kernel.shape)     # torch.Size([1, 3, 3, 3])

conv_func_img = torch.nn.functional.conv2d(
    img,
    kernel,     # weight 也就是卷积核
    bias=None   # 偏置
)
print("卷积后的形状", conv_func_img.shape)
T.ToPILImage()(conv_func_img).show()

# ——————调用类——————————————————————————————————————————————————————————————————
conv2d = torch.nn.Conv2d(
    3,     # 输入通道数
    1,  # 输出通道数
    3   # 卷积核大小
)
conv_img = conv2d(img)
print("卷积后的形状", conv_img.shape)
T.ToPILImage()(conv_img).show()
